#include "controller/dji_controller.h"

namespace Planner
{
    DJIController::DJIController(const ros::NodeHandle &node_handle)
        : nh_(node_handle)
    {
        init();
    }

    void DJIController::init()
    {
        init_gps_ = false;
        init_transform_ = false;
        prev_yaw_ = 0;
        ei_pos.setZero();
        ei_vel.setZero();
        ei_acc.setZero();

        C_EARTH = 6378137.0;
        nh_.param("dji_controller/gra", gra_, 9.8);
        nh_.param("dji_controller/mass", mass_, 6.3);
        nh_.param("dji_controller/max_thrust", max_thrust_, 212.9);
        nh_.param("dji_controller/vel_tol", tol_, 0.05);
        nh_.param("dji_controller/gps_or_laser", gps_or_laser_, false);

        nh_.param("dji_controller/kp_pos", kp_pos, 1.0);
        nh_.param("dji_controller/kp_vel", kp_vel, 1.0);
        nh_.param("dji_controller/kp_acc", kp_acc, 0.0);
        nh_.param("dji_controller/ki_pos", ki_pos, 0.1);
        nh_.param("dji_controller/ki_vel", ki_vel, 0.1);
        nh_.param("dji_controller/ki_acc", ki_acc, 0.0);

        string thrust_control_topic;
        nh_.param("dji_controller/thrust_control_topic",
                  thrust_control_topic, string("/dji_sdk/thrust_control"));

        control_command_pub_ = nh_.advertise<quad_msgs::ThrustAndAttitude>(
            thrust_control_topic, 10);
        curr_odom_pub_ = nh_.advertise<nav_msgs::Odometry>(
            "curr_odom_vis", 10);
        control_command_sub_ = nh_.subscribe("/control_command", 1,
                                             &DJIController::controlCommandCallback, this);
        gps_position_sub_ = nh_.subscribe("/dji_osdk_ros/gps_position", 1,
                                          &DJIController::gpsPositionCallback, this);
        gps_velocity_sub_ = nh_.subscribe("/dji_osdk_ros/velocity", 1,
                                          &DJIController::gpsVelocityCallback, this);
        dji_acceleration_sub_ = nh_.subscribe("/dji_osdk_ros/acceleration_ground_fused", 1,
                                              &DJIController::djiAccelerationCallback, this);
        dji_orien_sub_ = nh_.subscribe("/dji_osdk_ros/attitude", 1,
                                       &DJIController::djiOrienCallback, this);

        string laser_odom_topic;
        nh_.param("dji_controller/laser_odom_topic",
                  laser_odom_topic, string("odom"));
        laser_position_sub_ = nh_.subscribe(laser_odom_topic, 1,
                                            &DJIController::laserPositionCallback, this);
    }

    void DJIController::controlCommandCallback(const quad_msgs::ControlCommand &msg)
    {
        // 获取目标状态
        geometry2EigenVector(msg.position, desire_pos_);
        geometry2EigenVector(msg.velocity, desire_vel_);
        geometry2EigenVector(msg.acceleration, desire_acc_);

        if (desire_vel_.norm() > tol_)
        {
            Quaterniond tmp_orien;
            direction2quaternion(desire_vel_, tmp_orien);
            desire_yaw_ = get_yaw_from_quaternion(tmp_orien);
            prev_yaw_ = desire_yaw_;
        }
        else
        {
            desire_yaw_ = prev_yaw_;
        }

        calcOutput();
    }

    void DJIController::gpsPositionCallback(const sensor_msgs::NavSatFix &msg)
    {
        // 计算结果为ENU坐标系
        if (!init_gps_)
        {
            init_latitude_ = msg.latitude;
            init_longitude_ = msg.longitude;
            init_hegiht_ = msg.altitude;
            init_gps_ = true;
            ROS_WARN("Initial height: %f", init_hegiht_);
        }
        else
        {
            gps_pos_[1] = toRad(msg.latitude - init_latitude_) * C_EARTH;
            gps_pos_[0] = toRad(msg.longitude - init_longitude_) * C_EARTH * cos(toRad(msg.latitude));
            gps_pos_[2] = msg.altitude - init_hegiht_;
        }
    }

    void DJIController::gpsVelocityCallback(const geometry_msgs::Vector3Stamped &msg)
    {
        // DJI_OSDK_ROS中velocity话题发布的速度信息坐标系为ENU
        geometry2EigenVector(msg.vector, gps_vel_);
        // gps_vel_[0] = msg.vector.y;
        // gps_vel_[1] = msg.vector.x;
        // gps_vel_[2] = msg.vector.z;
    }

    void DJIController::djiAccelerationCallback(const geometry_msgs::Vector3Stamped &msg)
    {
        geometry2EigenVector(msg.vector, dji_acc_);
    }

    void DJIController::djiOrienCallback(const geometry_msgs::QuaternionStamped &msg)
    {
        // 姿态角为FLU坐标系相对于ENU的初始姿态
        geometry2EigenQuaternion(msg.quaternion, dji_orien_);

        if (!init_transform_)
        {
            flu2enu_ = dji_orien_;
            init_transform_ = true;
        }
    }

    void DJIController::laserPositionCallback(const nav_msgs::Odometry &msg)
    {
        geometry2EigenVector(msg.pose.pose.position, laser_pos_);
        geometry2EigenQuaternion(msg.pose.pose.orientation, laser_orien_);
    }

    // PID控制
    void DJIController::calcOutput()
    {
        // 确定里程计信息来源
        Vector3d curr_pos, curr_vel, curr_acc;
        Quaterniond curr_orien;
        double curr_yaw;
        if (gps_or_laser_)
        {
            curr_pos = laser_pos_;
            // SLAM无法获取速度、加速度信息
            curr_vel = desire_vel_;
            curr_acc = desire_acc_;
            curr_yaw = get_yaw_from_quaternion(laser_orien_);
            curr_orien = laser_orien_;
        }
        else
        {
            // TODO 此处为ENU坐标系下，需要转换到初始时刻的FLU坐标系下
            // curr_pos = gps_pos_;
            // curr_vel = gps_vel_;
            // curr_acc = dji_acc_;
            // curr_yaw = get_yaw_from_quaternion(dji_orien_);
            // curr_orien = dji_orien_;

            curr_pos = flu2enu_.inverse() * gps_pos_;
            curr_vel = flu2enu_.inverse() * gps_vel_;
            curr_orien = flu2enu_.inverse() * dji_orien_;
            curr_yaw = get_yaw_from_quaternion(curr_orien);
            ROS_INFO("before transform: (%f, %f, %f)", gps_pos_[0], gps_pos_[1], gps_pos_[2]);
        }
        ROS_INFO("desired pos: (%f, %f, %f), vel: (%f, %f, %f), acc: (%f, %f, %f)",
                 desire_pos_[0], desire_pos_[1], desire_pos_[2],
                 desire_vel_[0], desire_vel_[1], desire_vel_[2],
                 desire_acc_[0], desire_acc_[1], desire_acc_[2]);

        ROS_INFO("current pos: (%f, %f, %f), vel: (%f, %f, %f), acc: (%f, %f, %f)",
                 curr_pos[0], curr_pos[1], curr_pos[2],
                 curr_vel[0], curr_vel[1], curr_vel[2],
                 curr_acc[0], curr_acc[1], curr_acc[2]);

        // 先不考虑yaw角的问题
        // curr_yaw = desire_yaw_ = 0;
        // curr_orien = Eigen::Quaterniond::FromTwoVectors(Eigen::Vector3d(1, 0, 0), Eigen::Vector3d(1, 0, 0));

        // 计算误差项
        Vector3d ep_pos, ep_vel, ep_acc;
        ep_pos = desire_pos_ - curr_pos;
        ep_vel = desire_vel_ - curr_vel;
        ep_acc = desire_acc_ - curr_acc;
        double limit = 0.5;
        for (size_t k = 0; k < 3; ++k)
        {
            if (std::fabs(ep_pos(k)) < 0.2)
            {
                ei_pos(k) += ep_pos(k) / 50.0;
                limit_range(ei_pos(k), limit);
            }
            if (std::fabs(ep_vel(k)) < 0.2)
            {
                ei_vel(k) += ep_vel(k) / 50.0;
                limit_range(ei_vel(k), limit);
            }
            if (std::fabs(ep_acc(k)) < 0.2)
            {
                ei_acc(k) += ep_acc(k) / 50.0;
                limit_range(ei_acc(k), limit);
            }
        }
        Vector3d error = kp_pos * ep_pos + kp_vel * ep_vel + kp_acc * ep_acc +
                         ki_pos * ep_pos + ki_vel * ep_vel + ki_acc * ep_acc;

        // 限制推力大小
        Vector3d des_thrust = mass_ * error + Vector3d(0, 0, mass_ * gra_) + mass_ * desire_acc_;
        if (des_thrust[2] < 0.75 * mass_ * gra_)
        {
            ROS_INFO("thrust too slow, des_thrust: %f", des_thrust[2] / mass_);
            des_thrust[2] = 0.75 * mass_ * gra_;
        }
        if (des_thrust[2] > 1.25 * mass_ * gra_)
        {
            ROS_INFO("thrust too high, des_thrust: %f", des_thrust[2] / mass_);
            des_thrust[2] = 1.25 * mass_ * gra_;
        }
        // M300的最大俯仰角度为25度
        double angle_limit = 35.0;
        if (std::fabs(des_thrust[0] / des_thrust[2]) > std::tan(toRad(angle_limit)))
        {
            ROS_INFO("fx too high: %f", toDeg(std::atan2(des_thrust[0], des_thrust[2])));
            des_thrust[0] = des_thrust[0] / std::fabs(des_thrust[0]) * des_thrust[2] * std::tan(toRad(25.0));
        }
        if (std::fabs(des_thrust[1] / des_thrust[2]) > std::tan(toRad(angle_limit)))
        {
            ROS_INFO("fy too high: %f", toDeg(std::atan2(des_thrust[1], des_thrust[2])));
            des_thrust[1] = des_thrust[1] / std::fabs(des_thrust[1]) * des_thrust[2] * std::tan(toRad(25.0));
        }

        // 初始FLU与当前FLU坐标系相关变换
        Matrix3d wRc = rotz(curr_yaw);
        Vector3d F_c = wRc.transpose() * des_thrust;
        Matrix3d wRb_odom = curr_orien.toRotationMatrix();
        Vector3d z_b_curr = wRb_odom.col(2);
        double u1 = des_thrust.dot(z_b_curr); // 在z方向上的投影之和
        double fx = F_c(0);
        double fy = F_c(1);
        double fz = F_c(2);

        // 再次对thrust和roll、pitch、yaw取值做限定
        // TODO 需要对roll、pitch、yaw的符号进行确认
        //      command中yaw是相对于GROUND FRAME的，desire_yaw_是相对于机体坐标系的
        double roll = std::atan2(-fy, fz);
        roll = toDeg(roll);
        limit_range(roll, angle_limit);

        double pitch = std::atan2(fx, fz);
        pitch = toDeg(-pitch);
        limit_range(pitch, angle_limit);

        double init_yaw = get_yaw_from_quaternion(flu2enu_);
        double yaw = M_PI / 2 - (desire_yaw_ + init_yaw);
        std::cout << init_yaw << " " << desire_yaw_ << std::endl;
        ROS_INFO("init_yaw: %f, desired_yaw: %f", init_yaw, desire_yaw_);
        while (yaw > M_PI)
            yaw -= (2 * M_PI);
        while (yaw < -M_PI)
            yaw += (2 * M_PI);
        yaw = toDeg(yaw);

        limit_range(u1, 0.75 * mass_ * gra_, 1.25 * mass_ * gra_);
        double thrust = u1 / max_thrust_ * 100.0;

        // 发送控制指令
        quad_msgs::ThrustAndAttitude msg;
        msg.thrust = thrust;
        msg.roll = roll;
        msg.pitch = pitch;
        // msg.roll = -pitch;
        // msg.pitch = roll;
        
        msg.yaw = yaw;
        ROS_INFO("thrust: %f, roll: %f, pitch: %f, yaw: %f", msg.thrust, msg.roll, msg.pitch, msg.yaw);
        control_command_pub_.publish(msg);

        nav_msgs::Odometry odom_msg;
        odom_msg.header.frame_id = "world";
        odom_msg.pose.pose.position.x = curr_pos[0];
        odom_msg.pose.pose.position.y = curr_pos[1];
        odom_msg.pose.pose.position.z = curr_pos[2];
        odom_msg.pose.pose.orientation.x = curr_orien.x();
        odom_msg.pose.pose.orientation.y = curr_orien.y();
        odom_msg.pose.pose.orientation.z = curr_orien.z();
        odom_msg.pose.pose.orientation.w = curr_orien.w();
        curr_odom_pub_.publish(odom_msg);
    }
} // namespace Planner